from fastapi import APIRouter, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import os
import tempfile
import pytesseract
from PIL import Image
import shutil
from typing import Optional
import subprocess
import sys
import platform

# 从配置文件加载Tesseract配置
def load_tesseract_config():
    try:
        config_file = os.path.join("secret_files", "tesseract_config.txt")
        if os.path.exists(config_file):
            with open(config_file, "r") as f:
                for line in f:
                    if line.startswith("TESSERACT_CMD="):
                        path = line[len("TESSERACT_CMD="):].strip()
                        if os.path.exists(path):
                            print(f"从配置文件加载Tesseract路径: {path}")
                            os.environ["TESSERACT_CMD"] = path
                            return path
    except Exception as e:
        print(f"加载Tesseract配置失败: {str(e)}", file=sys.stderr)
    return None

# 优先从配置文件加载
config_path = load_tesseract_config()

# 智能检测Tesseract路径 - 基于系统类型和常见安装位置
def find_tesseract_path():
    # 如果从配置文件加载了有效路径，直接使用
    if config_path and os.path.exists(config_path):
        return config_path
        
    # 首先检查环境变量中是否已设置
    tesseract_cmd = os.environ.get('TESSERACT_CMD')
    if tesseract_cmd and os.path.exists(tesseract_cmd):
        return tesseract_cmd

    # 尝试直接运行tesseract命令
    try:
        subprocess.run(['tesseract', '--version'], 
                      stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                      text=True, timeout=2)
        return 'tesseract'  # 命令可用，使用默认命令
    except:
        pass  # 命令不可用，继续检查其他位置

    # 基于操作系统检查常见安装位置
    if platform.system() == 'Windows':
        possible_paths = [
            r'C:\Program Files\Tesseract-OCR\tesseract.exe',
            r'C:\Program Files (x86)\Tesseract-OCR\tesseract.exe',
            r'C:\Tesseract-OCR\tesseract.exe',
            # 添加D盘常见路径
            r'D:\Program Files\Tesseract-OCR\tesseract.exe',
            r'D:\Program Files (x86)\Tesseract-OCR\tesseract.exe',
            r'D:\Tesseract-OCR\tesseract.exe',
            r'D:\Software\Tesseract-OCR\tesseract.exe',
            r'D:\Tools\Tesseract-OCR\tesseract.exe',
            # 用户目录路径
            os.path.join(os.environ.get('LOCALAPPDATA', ''), 'Tesseract-OCR', 'tesseract.exe'),
            os.path.join(os.environ.get('APPDATA', ''), 'Tesseract-OCR', 'tesseract.exe'),
            # 添加用户环境变量PATH中的可能路径
            *[os.path.join(path, 'tesseract.exe') for path in os.environ.get('PATH', '').split(';') if 'tesseract' in path.lower()]
        ]
        
        # 在D盘根目录下搜索可能的Tesseract安装
        try:
            for root, dirs, files in os.walk('D:\\', topdown=True, onerror=None, followlinks=False):
                # 限制搜索深度以提高效率
                if root.count(os.sep) > 3:  # 不要搜索太深的目录
                    dirs.clear()  # 不再递归到子目录
                    continue
                
                # 排除某些系统目录
                dirs[:] = [d for d in dirs if d.lower() not in ['windows', '$recycle.bin', 'system volume information']]
                
                # 检查当前目录是否可能包含Tesseract
                if 'tesseract' in root.lower() and 'tesseract.exe' in [f.lower() for f in files]:
                    tesseract_exe = os.path.join(root, 'tesseract.exe')
                    if os.path.exists(tesseract_exe):
                        print(f"在D盘找到Tesseract: {tesseract_exe}")
                        return tesseract_exe
        except Exception as e:
            print(f"搜索D盘时出错: {str(e)}", file=sys.stderr)
    else:  # Linux/Mac
        possible_paths = [
            '/usr/bin/tesseract',
            '/usr/local/bin/tesseract',
            '/opt/local/bin/tesseract',
            '/opt/homebrew/bin/tesseract'
        ]

    # 检查路径是否存在
    for path in possible_paths:
        if os.path.exists(path):
            return path

    # 默认返回标准路径
    if platform.system() == 'Windows':
        return r'C:\Program Files\Tesseract-OCR\tesseract.exe'
    else:
        return 'tesseract'

# 设置Tesseract OCR路径
tesseract_path = find_tesseract_path()
pytesseract.pytesseract.tesseract_cmd = tesseract_path
print(f"使用Tesseract路径: {tesseract_path}")

# 尝试自动设置TESSDATA_PREFIX环境变量
def set_tessdata_prefix():
    # 如果已设置，不再处理
    if os.environ.get('TESSDATA_PREFIX'):
        return os.environ.get('TESSDATA_PREFIX')
        
    if os.path.exists(tesseract_path):
        # 根据tesseract.exe的路径推断tessdata目录
        possible_tessdata_dirs = [
            # 通常在tesseract-ocr目录下有tessdata子目录
            os.path.join(os.path.dirname(tesseract_path), 'tessdata'),
            os.path.join(os.path.dirname(os.path.dirname(tesseract_path)), 'tessdata'),
            # 或者在share/tessdata目录下
            os.path.join(os.path.dirname(os.path.dirname(tesseract_path)), 'share', 'tessdata'),
        ]
        
        # 检查可能的tessdata目录
        for tessdata_dir in possible_tessdata_dirs:
            if os.path.exists(tessdata_dir) and os.path.isdir(tessdata_dir):
                # 检查是否包含语言数据文件
                lang_files = [f for f in os.listdir(tessdata_dir) if f.endswith('.traineddata')]
                if lang_files:
                    os.environ['TESSDATA_PREFIX'] = tessdata_dir
                    print(f"设置TESSDATA_PREFIX={tessdata_dir}")
                    return tessdata_dir
                    
        # 如果没有找到，尝试在D盘搜索
        try:
            for root, dirs, files in os.walk('D:\\', topdown=True, onerror=None, followlinks=False):
                # 限制搜索深度
                if root.count(os.sep) > 4:
                    dirs.clear()
                    continue
                    
                # 看是否有tessdata目录
                if 'tessdata' in dirs:
                    tessdata_dir = os.path.join(root, 'tessdata')
                    # 检查内容
                    lang_files = [f for f in os.listdir(tessdata_dir) if f.endswith('.traineddata')]
                    if lang_files:
                        os.environ['TESSDATA_PREFIX'] = tessdata_dir
                        print(f"在D盘找到tessdata目录并设置TESSDATA_PREFIX={tessdata_dir}")
                        return tessdata_dir
        except Exception as e:
            print(f"搜索tessdata目录时出错: {str(e)}", file=sys.stderr)
            
    return None

# 尝试设置tessdata目录
tessdata_dir = set_tessdata_prefix()

# 检测Tesseract是否可用的函数
def is_tesseract_available():
    try:
        # 尝试执行tesseract命令
        result = subprocess.run([pytesseract.pytesseract.tesseract_cmd, "--version"], 
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                               text=True, timeout=5)
        if result.returncode == 0:
            version_text = result.stdout.strip().split('\n')[0]
            print(f"Tesseract 版本信息: {version_text}")
            return True
        return False
    except Exception as e:
        print(f"Tesseract OCR 检测失败: {str(e)}", file=sys.stderr)
        return False

# 检查Tesseract OCR是否可用
TESSERACT_AVAILABLE = is_tesseract_available()
if not TESSERACT_AVAILABLE:
    print(f"""
警告: Tesseract OCR 未安装或未在PATH中!
尝试使用的路径: {pytesseract.pytesseract.tesseract_cmd}

解决方法:
1. 确认Tesseract OCR已安装 (使用命令行运行'tesseract --version'检查)
2. 确保安装中文语言包
3. 将安装目录添加到系统PATH
4. 或者设置环境变量 TESSERACT_CMD 指向tesseract.exe的完整路径
5. 或者在代码中修改pytesseract.pytesseract.tesseract_cmd指向正确的路径

图像OCR功能将不可用，直到Tesseract被正确安装。
""", file=sys.stderr)

router = APIRouter(prefix="/api/pdfminer/image", tags=["图像文本处理"])

# 确保输入和输出目录存在
INPUT_DIR = "input_files"
OUTPUT_DIR = "output_files"
os.makedirs(INPUT_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 支持的图像文件类型
SUPPORTED_IMAGE_TYPES = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.gif']

@router.post("/extract")
async def extract_image_text(file: UploadFile = File(...), save_output: bool = True, lang: str = "chi_sim+eng"):
    """
    从上传的图像中提取文本
    - file: 上传的图像文件
    - save_output: 是否将结果保存到文件中(默认为true)
    - lang: OCR语言模型，默认为简体中文+英文(chi_sim+eng)
    """
    # 首先检查Tesseract是否可用
    if not TESSERACT_AVAILABLE:
        raise HTTPException(
            status_code=500,
            detail="Tesseract OCR未安装或未配置正确。请参考README文件安装Tesseract OCR，或联系管理员。"
        )
    
    # 检查文件类型
    file_ext = os.path.splitext(file.filename)[1].lower()
    if file_ext not in SUPPORTED_IMAGE_TYPES:
        raise HTTPException(
            status_code=400, 
            detail=f"不支持的文件类型，仅支持: {', '.join(SUPPORTED_IMAGE_TYPES)}"
        )
    
    # 保存上传的文件到临时文件
    temp_dir = tempfile.mkdtemp()
    temp_file_path = os.path.join(temp_dir, file.filename)
    
    try:
        with open(temp_file_path, "wb") as temp_file:
            content = await file.read()
            temp_file.write(content)
        
        # 保存到输入目录
        input_file_path = os.path.join(INPUT_DIR, file.filename)
        shutil.copy2(temp_file_path, input_file_path)
        
        # 使用Tesseract OCR提取文本
        try:
            image = Image.open(temp_file_path)
            text = pytesseract.image_to_string(image, lang=lang)
        except Exception as e:
            error_msg = str(e)
            if "Invalid language" in error_msg:
                raise HTTPException(
                    status_code=400, 
                    detail=f"无效的语言参数'{lang}'。可能未安装此语言包，请检查Tesseract语言设置。"
                )
            else:
                raise HTTPException(status_code=500, detail=f"OCR处理出错: {error_msg}")
        
        # 如果需要保存到文件
        if save_output:
            output_filename = os.path.splitext(file.filename)[0] + ".txt"
            output_path = os.path.join(OUTPUT_DIR, output_filename)
            with open(output_path, "w", encoding="utf-8") as out_file:
                out_file.write(text)
            
            # 获取绝对路径
            abs_output_path = os.path.abspath(output_path)
            
            return JSONResponse(content={
                "message": "图像文本提取成功",
                "filename": file.filename,
                "output_file": output_filename,
                "output_path": output_path,
                "absolute_path": abs_output_path,
                "text_preview": text[:500] + "..." if len(text) > 500 else text
            })
        
        return JSONResponse(content={
            "message": "图像文本提取成功",
            "filename": file.filename,
            "text": text
        })
    
    except Exception as e:
        if isinstance(e, HTTPException):
            raise e
        raise HTTPException(status_code=500, detail=f"处理图像时出错: {str(e)}")
    
    finally:
        # 清理临时文件
        shutil.rmtree(temp_dir, ignore_errors=True)

@router.get("/files")
async def list_image_files():
    """列出输入目录中的所有图像文件"""
    image_files = [
        f for f in os.listdir(INPUT_DIR) 
        if os.path.splitext(f)[1].lower() in SUPPORTED_IMAGE_TYPES
    ]
    return {"image_files": image_files}

@router.post("/batch_extract")
async def batch_extract_image_text(files: list[UploadFile] = File(...), save_output: bool = True, lang: str = "chi_sim+eng"):
    """
    批量处理多张图片
    - files: 多个图像文件
    - save_output: 是否将结果保存到文件中(默认为true)
    - lang: OCR语言模型，默认为简体中文+英文(chi_sim+eng)
    """
    # 首先检查Tesseract是否可用
    if not TESSERACT_AVAILABLE:
        raise HTTPException(
            status_code=500,
            detail="Tesseract OCR未安装或未配置正确。请参考README文件安装Tesseract OCR，或联系管理员。"
        )
    
    if len(files) == 0:
        raise HTTPException(status_code=400, detail="未提供文件")
    
    results = []
    
    for file in files:
        # 检查文件类型
        file_ext = os.path.splitext(file.filename)[1].lower()
        if file_ext not in SUPPORTED_IMAGE_TYPES:
            results.append({
                "filename": file.filename,
                "status": "error",
                "message": f"不支持的文件类型，仅支持: {', '.join(SUPPORTED_IMAGE_TYPES)}"
            })
            continue
        
        # 保存上传的文件到临时文件
        temp_dir = tempfile.mkdtemp()
        temp_file_path = os.path.join(temp_dir, file.filename)
        
        try:
            with open(temp_file_path, "wb") as temp_file:
                content = await file.read()
                temp_file.write(content)
            
            # 保存到输入目录
            input_file_path = os.path.join(INPUT_DIR, file.filename)
            shutil.copy2(temp_file_path, input_file_path)
            
            # 使用Tesseract OCR提取文本
            try:
                image = Image.open(temp_file_path)
                text = pytesseract.image_to_string(image, lang=lang)
            except Exception as e:
                error_msg = str(e)
                if "Invalid language" in error_msg:
                    results.append({
                        "filename": file.filename,
                        "status": "error",
                        "message": f"无效的语言参数'{lang}'。可能未安装此语言包。"
                    })
                else:
                    results.append({
                        "filename": file.filename,
                        "status": "error",
                        "message": f"OCR处理出错: {error_msg}"
                    })
                continue
            
            # 如果需要保存到文件
            if save_output:
                output_filename = os.path.splitext(file.filename)[0] + ".txt"
                output_path = os.path.join(OUTPUT_DIR, output_filename)
                with open(output_path, "w", encoding="utf-8") as out_file:
                    out_file.write(text)
                
                # 获取绝对路径
                abs_output_path = os.path.abspath(output_path)
                
                results.append({
                    "filename": file.filename,
                    "status": "success",
                    "output_file": output_filename,
                    "output_path": output_path,
                    "absolute_path": abs_output_path,
                    "text_preview": text[:200] + "..." if len(text) > 200 else text
                })
            else:
                results.append({
                    "filename": file.filename,
                    "status": "success",
                    "text": text
                })
        
        except Exception as e:
            results.append({
                "filename": file.filename,
                "status": "error",
                "message": f"处理图像时出错: {str(e)}"
            })
        
        finally:
            # 清理临时文件
            shutil.rmtree(temp_dir, ignore_errors=True)
    
    return JSONResponse(content={
        "message": f"处理完成 {len(results)} 个文件",
        "results": results
    })

@router.get("/tesseract-status")
async def tesseract_status():
    """获取Tesseract OCR的安装状态和系统信息，用于诊断问题"""
    # 系统信息
    system_info = {
        "os": platform.system(),
        "os_version": platform.version(),
        "python_version": platform.python_version(),
        "platform": platform.platform(),
        "architecture": platform.architecture(),
        "processor": platform.processor()
    }
    
    # Tesseract 信息
    tesseract_info = {
        "path_used": pytesseract.pytesseract.tesseract_cmd,
        "path_exists": os.path.exists(pytesseract.pytesseract.tesseract_cmd),
        "is_available": TESSERACT_AVAILABLE,
        "tessdata_prefix": os.environ.get("TESSDATA_PREFIX", "未设置"),
        "d_drive_exists": os.path.exists("D:\\")
    }
    
    # 尝试获取版本信息
    version_info = {}
    try:
        result = subprocess.run([pytesseract.pytesseract.tesseract_cmd, "--version"], 
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                               text=True, timeout=5)
        if result.returncode == 0:
            version_lines = result.stdout.strip().split('\n')
            version_info["version"] = version_lines[0] if version_lines else "未知"
            version_info["full_output"] = result.stdout
        else:
            version_info["error"] = result.stderr
    except Exception as e:
        version_info["error"] = str(e)
    
    # 环境变量信息
    env_info = {
        "PATH": os.environ.get("PATH", ""),
        "TESSERACT_CMD": os.environ.get("TESSERACT_CMD", "未设置"),
        "TESSDATA_PREFIX": os.environ.get("TESSDATA_PREFIX", "未设置")
    }
    
    # 搜索可用语言
    lang_info = {"available_languages": []}
    if TESSERACT_AVAILABLE:
        try:
            result = subprocess.run([pytesseract.pytesseract.tesseract_cmd, "--list-langs"], 
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                text=True, timeout=5)
            if result.returncode == 0:
                # 通常语言列表在stderr
                langs = result.stderr.strip().split('\n')
                if len(langs) > 1:  # 第一行通常是提示信息
                    lang_info["available_languages"] = langs[1:]
            else:
                lang_info["error"] = result.stderr
        except Exception as e:
            lang_info["error"] = str(e)
    
    # D盘Tesseract搜索
    d_drive_info = {"searched": False, "found_paths": []}
    if os.path.exists("D:\\") and not TESSERACT_AVAILABLE:
        d_drive_info["searched"] = True
        try:
            for root, dirs, files in os.walk('D:\\', topdown=True, onerror=None, followlinks=False):
                if root.count(os.sep) > 3:
                    dirs.clear()
                    continue
                
                # 排除系统目录    
                dirs[:] = [d for d in dirs if d.lower() not in ['windows', '$recycle.bin', 'system volume information']]
                
                # 搜索tesseract.exe
                if any('tesseract' in f.lower() for f in files):
                    for f in files:
                        if 'tesseract' in f.lower() and f.lower().endswith('.exe'):
                            full_path = os.path.join(root, f)
                            d_drive_info["found_paths"].append(full_path)
        except Exception as e:
            d_drive_info["error"] = str(e)
    
    # 返回所有收集的信息
    return {
        "system_info": system_info,
        "tesseract_info": tesseract_info,
        "version_info": version_info,
        "env_info": env_info,
        "lang_info": lang_info,
        "d_drive_info": d_drive_info,
        "diagnose": {
            "likely_issue": "找不到Tesseract可执行文件" if not tesseract_info["path_exists"] else 
                           "Tesseract路径配置错误" if not tesseract_info["is_available"] else
                           "Tesseract正常工作",
            "suggestions": [
                "确认Tesseract已安装且在PATH中" if not tesseract_info["is_available"] else "Tesseract配置正确",
                "检查语言包是否已安装" if tesseract_info["is_available"] and "chi_sim" not in str(lang_info["available_languages"]) else "中文语言包检查",
                f"考虑手动设置路径: pytesseract.pytesseract.tesseract_cmd = r'{os.path.join(os.path.expanduser('~'), 'AppData', 'Local', 'Programs', 'Tesseract-OCR', 'tesseract.exe')}'" if platform.system() == "Windows" and not tesseract_info["is_available"] else "",
                *([f"在D盘找到Tesseract: {path}，可以调用 /set-tesseract-path 设置" for path in d_drive_info.get("found_paths", [])] if d_drive_info.get("found_paths") else [])
            ]
        }
    }

@router.post("/set-tesseract-path")
async def set_tesseract_path(path: str):
    """设置Tesseract OCR可执行文件的路径"""
    global TESSERACT_AVAILABLE
    
    # 验证路径是否存在
    if not os.path.exists(path):
        return {
            "success": False,
            "message": f"指定的路径不存在: {path}",
            "status": "error"
        }
    
    # 尝试使用这个路径
    old_path = pytesseract.pytesseract.tesseract_cmd
    pytesseract.pytesseract.tesseract_cmd = path
    
    # 检查新路径是否有效
    is_valid = False
    version_info = ""
    try:
        result = subprocess.run([path, "--version"], 
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                               text=True, timeout=5)
        if result.returncode == 0:
            is_valid = True
            version_info = result.stdout.strip()
            TESSERACT_AVAILABLE = True
    except Exception as e:
        # 如果失败，恢复原来的路径
        pytesseract.pytesseract.tesseract_cmd = old_path
        return {
            "success": False,
            "message": f"路径无效，无法执行Tesseract: {str(e)}",
            "status": "error"
        }
    
    if is_valid:
        # 保存到环境变量，使重启后仍然有效
        os.environ["TESSERACT_CMD"] = path
        
        # 创建一个配置文件保存路径（可选）
        try:
            config_dir = "secret_files"
            os.makedirs(config_dir, exist_ok=True)
            with open(os.path.join(config_dir, "tesseract_config.txt"), "w") as f:
                f.write(f"TESSERACT_CMD={path}")
        except:
            pass  # 忽略配置文件创建失败
            
        return {
            "success": True,
            "message": f"Tesseract路径设置成功: {path}",
            "version": version_info,
            "status": "success"
        }
    else:
        # 恢复原来的路径
        pytesseract.pytesseract.tesseract_cmd = old_path
        return {
            "success": False,
            "message": "路径无效，无法执行Tesseract",
            "status": "error"
        } 